// Communication Module
// SendMessage and PostMessage is used for communication
// 
#include "stdafx.h"
#include "communication.h"
#include "SharedMemory.h"
#include "SocketHelper.h"
#include <stdlib.h>
#include <strsafe.h>
#include <WinCrypt.h>

#pragma data_seg("Comm")
HWND g_MainAppMsgHwnd = NULL;
#pragma data_seg()

extern DWORD m_currentProcessID;

BOOL m_bIsConsoleCreated = FALSE;
HANDLE m_console = NULL;

//////////////////////////////////////////////////////////////////////////
// Exports
//////////////////////////////////////////////////////////////////////////

VOID WINAPI SetMainAppHwnd(HWND hWnd)
{
	g_MainAppMsgHwnd = hWnd;
	// SayHelloToMainApp();
}

VOID WINAPI Test(ULONG slong)
{
	SayHelloToMainApp();
}

//////////////////////////////////////////////////////////////////////////
// Used inside
//////////////////////////////////////////////////////////////////////////

DWORD GetRandomInteger()
{
	HCRYPTPROV hCryptProv = NULL;
	DWORD dwRand;
	CryptAcquireContext((HCRYPTPROV*)&hCryptProv, NULL, NULL, PROV_RSA_FULL, 0); 
	CryptGenRandom(hCryptProv, sizeof(DWORD), (BYTE*)&dwRand);
	CryptReleaseContext(hCryptProv, 0); 
	return dwRand;
}

INT SayHelloToMainApp()
{
	if(g_MainAppMsgHwnd != NULL)
	{
		// ::MessageBox(NULL, _T("Sending Hello!"), _T("..."), MB_OK | MB_ICONINFORMATION);
		return SendMessage(g_MainAppMsgHwnd, MSG_HELLO, m_currentProcessID, 0);
	}
	else
	{
/*		::MessageBox(NULL, _T("g_MainAppMsgHwnd is null."), _T("..."), MB_OK | MB_ICONINFORMATION);*/
		return NULLMSGHWND;
	}
}

BOOL Filter_connect(SOCKET s, const struct sockaddr FAR * name, int namelen)
{
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;
		TCHAR strBuffer[20] = {0};

		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		msg.hookType = WHT_CONNECT;
		msg.lpData = NULL;
		msg.sockRelated = s;
		msg.uDataSharedMemoryName = 0;
		msg.nFlags = 0;
		msg.uDataLength = 0;
		memcpy_s(&msg.socketAddr, sizeof(struct sockaddr), name, sizeof(struct sockaddr));
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		return SendMessage(g_MainAppMsgHwnd, MSG_FILTER_CONNECT, uSharedMemoryName, m_currentProcessID);
	}
	return TRUE;
}

BOOL Filter_send(SOCKET s, const char FAR *buf, int len, int flags)
{
	INT ret;
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;
		CSharedMemory *dataSharedMemory = new CSharedMemory;
		TCHAR strBuffer[20] = {0};
		
		TCHAR strDataSharedMemoryName[20] = {0};
		DWORD uDataSharedMemoryName = GetRandomInteger();
		wsprintf(strDataSharedMemoryName, _T("%u"), uDataSharedMemoryName);
		dataSharedMemory->Create(strDataSharedMemoryName, len);
		dataSharedMemory->Write(buf, len);

		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		msg.hookType = WHT_SEND;
		msg.lpData = NULL;
		msg.sockRelated = s;
		msg.uDataSharedMemoryName = uDataSharedMemoryName;
		msg.nFlags = flags;
		msg.uDataLength = len;
		ZeroMemory(&msg.socketAddr, sizeof(struct sockaddr));
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		ret = SendMessage(g_MainAppMsgHwnd, MSG_FILTER_SEND, uSharedMemoryName, m_currentProcessID);
		delete shareMemory;
		delete dataSharedMemory;
	}
	return ret;
}

BOOL Filter_sendto(SOCKET s, const char FAR *buf, int len, int flags, const struct sockaddr* to, int tolen)
{
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;
		TCHAR strBuffer[20] = {0};
		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		msg.hookType = WHT_SENDTO;
		msg.lpData = NULL;
		msg.sockRelated = s;
		msg.nFlags = flags;
		msg.uDataLength = len;
		memcpy_s(&msg.socketAddr, sizeof(struct sockaddr), to, sizeof(struct sockaddr));
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		return SendMessage(g_MainAppMsgHwnd, MSG_FILTER_SEND, uSharedMemoryName, m_currentProcessID);
	}
	return TRUE;
}

BOOL Filter_recv(SOCKET s, char FAR *buf, int len, int flags)
{
/*	CHAR buffer[100] = {0};*/
	if(g_MainAppMsgHwnd != NULL)
	{
// 		sprintf(buffer, "enter len = %d, thread = %x, WSAGetLastError = %x \n", len, GetCurrentThreadId(), WSAGetLastError());
// 		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);

		CSharedMemory *shareMemory = new CSharedMemory;

 		CSharedMemory *dataSharedMemory = new CSharedMemory;
  		TCHAR strDataSharedMemoryName[20] = {0};
  		DWORD uDataSharedMemoryName = GetRandomInteger();
  		wsprintf(strDataSharedMemoryName, _T("%u"), uDataSharedMemoryName);
  		dataSharedMemory->Create(strDataSharedMemoryName, len);
		dataSharedMemory->Write(buf, len);

 		TCHAR strBuffer[20] = {0};
 		int uSharedMemoryName = GetRandomInteger();
 		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
 		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
 		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
 		msg.hookType = WHT_RECV;
 		msg.lpData = NULL;
 		msg.sockRelated = s;
		msg.uDataSharedMemoryName = uDataSharedMemoryName;
 		msg.nFlags = flags;
 		msg.uDataLength = len;
		ZeroMemory(&msg.socketAddr, sizeof(struct sockaddr));
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));

 		LRESULT ret = SendMessage(g_MainAppMsgHwnd, MSG_FILTER_RECV, uSharedMemoryName, m_currentProcessID);

 		delete shareMemory;
 		delete dataSharedMemory;
//  		sprintf(buffer, "leave thread = %x\n", GetCurrentThreadId());
// 		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
		return TRUE;
	}
	return TRUE;
}

BOOL Filter_recvfrom(SOCKET s, char* buf, int len, int flags, struct sockaddr* from, int* fromlen)
{
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;

		CSharedMemory *dataSharedMemory = new CSharedMemory;
		TCHAR strDataSharedMemoryName[20] = {0};
		DWORD uDataSharedMemoryName = GetRandomInteger();
		wsprintf(strDataSharedMemoryName, _T("%u"), uDataSharedMemoryName);
		dataSharedMemory->Create(strDataSharedMemoryName, len);
		dataSharedMemory->Write(buf, len);

		TCHAR strBuffer[20] = {0};
		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		msg.hookType = WHT_RECVFROM;
		msg.lpData = NULL;
		msg.uDataSharedMemoryName = uDataSharedMemoryName;
		msg.sockRelated = s;
		msg.nFlags = flags;
		msg.uDataLength = len;
		memcpy(&msg.socketAddr, from, *fromlen);
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		return SendMessage(g_MainAppMsgHwnd, MSG_FILTER_RECVFROM, uSharedMemoryName, m_currentProcessID);
	}
	return TRUE;
}

BOOL Filter_WSASend(SOCKET s, LPWSABUF lpBuffers, 
	DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, 
	LPWSAOVERLAPPED lpOverlapped, 
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
	INT ret;
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;
		CSharedMemory *dataSharedMemory = new CSharedMemory;
		TCHAR strBuffer[20] = {0};

		DWORD nBufferSize = 0;
		for(int i = 0; i < dwBufferCount; ++i)
		{
			nBufferSize += lpBuffers[i].len;
		}
		TCHAR strDataSharedMemoryName[20] = {0};
		DWORD uDataSharedMemoryName = GetRandomInteger();
		wsprintf(strDataSharedMemoryName, _T("%u"), uDataSharedMemoryName);
		dataSharedMemory->Create(strDataSharedMemoryName, nBufferSize);

		DWORD nSizeWritten = 0;
		for(int i = 0; i < dwBufferCount && nSizeWritten < nBufferSize; ++i)
		{
			DWORD nBytesToWrite = min(lpBuffers[i].len, nBufferSize - nSizeWritten);
			dataSharedMemory->Write(lpBuffers[i].buf, nBytesToWrite);
			nSizeWritten += nBytesToWrite;
		}

		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		msg.hookType = WHT_WSASEND;
		msg.lpData = NULL;
		msg.sockRelated = s;
		msg.uDataSharedMemoryName = uDataSharedMemoryName;
		msg.nFlags = dwFlags;
		msg.uDataLength = nBufferSize;
		msg.lpOverlapped = lpOverlapped;
		if(lpOverlapped != NULL)
		{
			memcpy_s(&msg.lpOverlapped, sizeof(WSAOVERLAPPED), lpOverlapped, sizeof(WSAOVERLAPPED));
		}
		ZeroMemory(&msg.socketAddr, sizeof(struct sockaddr));
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		ret = SendMessage(g_MainAppMsgHwnd, MSG_FILTER_WSASEND, uSharedMemoryName, m_currentProcessID);

		delete shareMemory;
		delete dataSharedMemory;
	}
	return ret;
}

BOOL Filter_WSARecv(
	SOCKET s,                                               
	LPWSABUF lpBuffers,                                     
	DWORD dwBufferCount,                                    
	LPDWORD lpNumberOfBytesRecvd,                           
	LPDWORD lpFlags,                                        
	LPWSAOVERLAPPED lpOverlapped,                           
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine  
	)
{
	INT ret = 0;
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;

		CSharedMemory *dataSharedMemory = new CSharedMemory;
		TCHAR strDataSharedMemoryName[20] = {0};
		DWORD uDataSharedMemoryName = GetRandomInteger();
		wsprintf(strDataSharedMemoryName, _T("%u"), uDataSharedMemoryName);
		dataSharedMemory->Create(strDataSharedMemoryName, *lpNumberOfBytesRecvd);
		dataSharedMemory->Write(lpBuffers->buf, *lpNumberOfBytesRecvd);

		TCHAR strBuffer[20] = {0};
		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		msg.hookType = WHT_WSARECV;
		msg.lpData = NULL;
		msg.uDataSharedMemoryName = uDataSharedMemoryName;
		msg.sockRelated = s;
		msg.nFlags = *lpFlags;
		msg.uDataLength = *lpNumberOfBytesRecvd;
		msg.lpOverlapped = lpOverlapped;
		memcpy_s(&msg.WSAOverlapped, sizeof(WSAOVERLAPPED), lpOverlapped, sizeof(WSAOVERLAPPED));
		msg.lpCompletionRoutine = (DWORD_PTR)lpCompletionRoutine;
		ZeroMemory(&msg.socketAddr, sizeof(struct sockaddr));
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		ret = SendMessage(g_MainAppMsgHwnd, MSG_FILTER_WSARECV, uSharedMemoryName, m_currentProcessID);
		delete dataSharedMemory;
		delete shareMemory;
	}
	return TRUE;
}

BOOL Filter_WSARecvFrom(
	SOCKET s,                                               
	LPWSABUF lpBuffers,                                     
	DWORD dwBufferCount,                                    
	LPDWORD lpNumberOfBytesRecvd,                           
	LPDWORD lpFlags,                                        
	struct sockaddr FAR *lpFrom,                            
	LPINT lpFromlen,                                        
	LPWSAOVERLAPPED lpOverlapped,                           
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine  
	)
{
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;

		CSharedMemory *dataSharedMemory = new CSharedMemory;
		TCHAR strDataSharedMemoryName[20] = {0};
		DWORD uDataSharedMemoryName = GetRandomInteger();
		wsprintf(strDataSharedMemoryName, _T("%u"), uDataSharedMemoryName);
		dataSharedMemory->Create(strDataSharedMemoryName, lpBuffers->len);
		dataSharedMemory->Write(lpBuffers->buf, lpBuffers->len);

		TCHAR strBuffer[20] = {0};
		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		msg.hookType = WHT_WSARECVFROM;
		msg.lpData = NULL;
		msg.uDataSharedMemoryName = uDataSharedMemoryName;
		msg.sockRelated = s;
		msg.nFlags = *lpFlags;
		msg.uDataLength = lpBuffers->len;
		memcpy_s(&msg.socketAddr, sizeof(struct sockaddr), lpFrom, *lpFromlen);
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		return SendMessage(g_MainAppMsgHwnd, MSG_FILTER_WSARECVFROM, uSharedMemoryName, m_currentProcessID);
	}
	return TRUE;
}

BOOL Notice_WSARecv(SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount,                                    
		LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags, LPWSAOVERLAPPED lpOverlapped,
		LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
	if(g_MainAppMsgHwnd != NULL)
	{
			CSharedMemory *shareMemory = new CSharedMemory;

			TCHAR strBuffer[20] = {0};
			int uSharedMemoryName = GetRandomInteger();
			wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
			shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
			WINSOCKHOOKMSG msg;
			ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
			msg.hookType = WHT_NOTICE_WSARECV;
			msg.lpData = NULL;
			msg.uDataSharedMemoryName = 0;
			msg.sockRelated = s;
			msg.nFlags = *lpFlags;
			msg.lpWSABuffers = lpBuffers;
			msg.dwBufferCount = dwBufferCount;
			msg.uDataLength = 0;
			msg.lpOverlapped = lpOverlapped;
			if(lpOverlapped != NULL)
			{
				memcpy_s(&msg.WSAOverlapped, sizeof(WSAOVERLAPPED), lpOverlapped, sizeof(WSAOVERLAPPED));
			}
			msg.lpCompletionRoutine = (DWORD_PTR)lpCompletionRoutine;
			ZeroMemory(&msg.socketAddr, sizeof(struct sockaddr));
			shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
			int ret = SendMessage(g_MainAppMsgHwnd, MSG_NOTICE_WSARECV, uSharedMemoryName, m_currentProcessID);
		
			delete shareMemory;
			
			return ret;
		}
		return TRUE;
	}

BOOL Notice_WSAGetOverlappedResult(SOCKET s, LPWSAOVERLAPPED lpOverlapped, 
	LPDWORD lpcbTransfer, BOOL fWait, LPDWORD lpdwFlags, CWSABufferInfo* pBufferInfo)
{
	if(g_MainAppMsgHwnd != NULL)
	{
		CSharedMemory *shareMemory = new CSharedMemory;

		DWORD nSize = 0;
		DWORD uDataSharedMemoryName = 0;
		CSharedMemory *dataSharedMemory = NULL;
		if(pBufferInfo != NULL)
		{
			dataSharedMemory = new CSharedMemory;
			TCHAR strDataSharedMemoryName[20] = {0};
			uDataSharedMemoryName = GetRandomInteger();
			wsprintf(strDataSharedMemoryName, _T("%u"), uDataSharedMemoryName);
			dataSharedMemory->Create(strDataSharedMemoryName, *lpcbTransfer);
			
			for(int i = 0; i < pBufferInfo->dwBufferCount && nSize < *lpcbTransfer; ++i)
			{
				INT nWriteLength = min(pBufferInfo->lpBuffers[i].len, *lpcbTransfer - nSize);
				dataSharedMemory->Write(pBufferInfo->lpBuffers[i].buf, nWriteLength);
				nSize += nWriteLength;
			}
		}

		TCHAR strBuffer[20] = {0};
		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(WINSOCKHOOKMSG));
		WINSOCKHOOKMSG msg;
		ZeroMemory(&msg, sizeof(WINSOCKHOOKMSG));
		if(pBufferInfo != NULL)
		{
			msg.hookType = pBufferInfo->originalHookType;
		}
		else
		{
			msg.hookType = WHT_NOTICE_WSAGETOVERLAPPEDRESULT;
		}
		msg.overlappedHookType = WHT_NOTICE_WSAGETOVERLAPPEDRESULT;
		msg.lpData = NULL;
		msg.uDataSharedMemoryName = uDataSharedMemoryName;
		msg.sockRelated = s;
		msg.nFlags = *lpdwFlags;
		msg.uDataLength = *lpcbTransfer;
		msg.lpOverlapped = lpOverlapped;
		memcpy_s(&msg.WSAOverlapped, sizeof(WSAOVERLAPPED), lpOverlapped, sizeof(WSAOVERLAPPED));
		ZeroMemory(&msg.socketAddr, sizeof(struct sockaddr));
		shareMemory->Write(&msg, sizeof(WINSOCKHOOKMSG));
		int ret = SendMessage(g_MainAppMsgHwnd, MSG_NOTICE_WSAGETOVERLAPPEDRESULT, uSharedMemoryName, m_currentProcessID);
	
		if(dataSharedMemory != NULL)
		{
			delete dataSharedMemory;
		}
		delete shareMemory;

		return ret;
	}
	return TRUE;
}

INT CheckAvailableCacheData(SOCKET s)
{
	if(g_MainAppMsgHwnd != NULL)
	{
		CacheDataStatus ret = (CacheDataStatus)SendMessage(g_MainAppMsgHwnd, MSG_CHECK_AVAILABLE_CACHE_DATA, (WPARAM)s, m_currentProcessID);
		return ret;
	}
	return CDS_NOTAVAILABLE;
}

INT GetAvailableCacheData(SOCKET s, PUCHAR lpData, INT nMaxSize)
{
	if(g_MainAppMsgHwnd != NULL)
	{
#ifdef CREATE_CONSOLE
		CHAR buffer[200] = {0};
		sprintf(buffer, "!!!!!!!!!!!!!!!!!!!!!!!!! Get available cache data!!!!!!!!!!!");
		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

		CSharedMemory *shareMemory = new CSharedMemory;
		TCHAR strBuffer[20] = {0};
		int uSharedMemoryName = GetRandomInteger();
		wsprintf(strBuffer, _T("%u"), uSharedMemoryName);
		shareMemory->Create(strBuffer, sizeof(CACHEDATAMSG));
		CACHEDATAMSG msg;
		msg.sockRelated = s;
		msg.nDataLength = nMaxSize;
		shareMemory->Write(&msg, sizeof(CACHEDATAMSG));
		INT nRecvDataStatus = nRecvDataStatus = SendMessage(g_MainAppMsgHwnd, MSG_GET_AVAILABLE_CACHE_DATA, uSharedMemoryName, m_currentProcessID);
		PCACHEDATAMSG pNewMsg = (PCACHEDATAMSG)shareMemory->GetBuffer();
		INT nDataSize = pNewMsg->nDataLength;
		if(nRecvDataStatus == DataReceived)
		{
			memcpy_s(lpData, nMaxSize, pNewMsg->lpData, nDataSize);
		}
		else if(nRecvDataStatus == EmptyBuffer)
		{
			/* More time is needed! */
			nDataSize = -1;
		}
		else /* No more data available */
		{
			nDataSize = 0;
		}
		delete shareMemory;
		return nDataSize;
	}
	return 0;
}

INT CreateProcessInform(INT nOldProcessID, int nNewProcessID)
{
	if(g_MainAppMsgHwnd != NULL)
	{
		SendMessage(g_MainAppMsgHwnd, MSG_NOTICE_CREATEPROCESSASUSER, nOldProcessID, nNewProcessID);
	}
	return 0;
}